"""
智能裁剪引擎
包含重叠裁剪算法和尺寸标准化处理
"""

import math
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Any, Tuple, Optional
from PIL import Image, ImageDraw
import logging
from pathlib import Path
# import numpy as np  # 暂时注释掉，如果需要可以重新启用
from datetime import datetime
import os

class CroppingEngine:
    """智能裁剪引擎"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化裁剪引擎"""
        self.config = config
        self.logger = logging.getLogger('image_cropping_system')
        
        # 目标尺寸配置
        self.target_sizes = config.get('target_sizes', {
            'square': [1024, 1024],
            'portrait': [768, 1024],
            'small': [512, 512]
        })
        
        # 裁剪配置
        self.overlap_ratio = config.get('processing', {}).get('overlap_ratio', 0.25)
        self.output_format = config.get('output', {}).get('format', 'PNG')
        self.output_quality = config.get('output', {}).get('quality', 95)
        self.naming_convention = config.get('output', {}).get('naming_convention', 'intelligent')
        
        # 性能优化配置
        self.max_workers = config.get('threading', {}).get('max_workers', 4)
        self.enable_parallel = config.get('processing', {}).get('enable_parallel', True)
        self.memory_optimization = config.get('processing', {}).get('memory_optimization', True)
        
        # 创建输出目录
        self.output_dir = Path('output/processed')
        for size_name in self.target_sizes.keys():
            (self.output_dir / size_name).mkdir(parents=True, exist_ok=True)
        
        # 强制使用PNG格式以支持透明度
        self.output_format = 'PNG'
        
        # 创建线程池
        self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
        
        # IO优化：预分配文件描述符
        if hasattr(os, 'posix_fadvise'):
            self.logger.info("启用IO优化：预分配文件描述符")
    
    async def crop_images(self, files: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量裁剪图片 - 支持并行处理"""
        if not files:
            return []
        
        self.logger.info(f"开始批量裁剪 {len(files)} 个文件，并行模式: {self.enable_parallel}")
        
        if self.enable_parallel and len(files) > 1:
            # 并行处理多个文件
            tasks = [self._crop_single_image_async(file_info) for file_info in files]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 处理异常结果
            processed_results = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    self.logger.error(f"文件 {files[i]['processed_name']} 处理失败: {result}")
                    processed_results.append({
                        **files[i],
                        'cropping_status': 'error',
                        'error': str(result),
                        'cropped_files': []
                    })
                else:
                    processed_results.extend(result)
            return processed_results
        else:
            # 串行处理
            results = []
            for file_info in files:
                try:
                    self.logger.info(f"开始裁剪图片: {file_info['processed_name']}")
                    cropped_results = await self._crop_single_image(file_info)
                    results.extend(cropped_results)
                except Exception as e:
                    self.logger.error(f"裁剪图片时发生错误 {file_info['processed_name']}: {e}")
                    results.append({
                        **file_info,
                        'cropping_status': 'error',
                        'error': str(e),
                        'cropped_files': []
                    })
            return results
    
    async def _crop_single_image_async(self, file_info: Dict[str, Any]) -> List[Dict[str, Any]]:
        """异步裁剪单张图片"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(self.executor, self._crop_single_image_sync, file_info)
    
    def _crop_single_image_sync(self, file_info: Dict[str, Any]) -> List[Dict[str, Any]]:
        """同步裁剪单张图片（在线程池中执行）"""
        try:
            with Image.open(file_info['processed_path']) as img:
                content_bbox = file_info['content_bbox'][:4]  # x1, y1, x2, y2
                confidence = file_info.get('confidence', 0.0)
                
                cropped_files = []
                
                # 为每个目标尺寸进行裁剪
                for size_name, target_size in self.target_sizes.items():
                    target_width, target_height = target_size
                    
                    # 检查是否需要重叠裁剪
                    if self._needs_overlap_cropping(content_bbox, target_size):
                        # 重叠裁剪
                        overlap_crops = self._smart_overlap_crop(
                            img, content_bbox, target_size, self.overlap_ratio
                        )
                        
                        for i, (crop_img, crop_coords) in enumerate(overlap_crops):
                            filename = self._generate_filename(
                                file_info, size_name, i, len(overlap_crops), crop_coords
                            )
                            
                            output_path = self.output_dir / size_name / filename
                            crop_img.save(output_path, format=self.output_format, quality=self.output_quality)
                            
                            cropped_files.append({
                                'filename': filename,
                                'path': str(output_path),
                                'size_name': size_name,
                                'width': target_width,
                                'height': target_height,
                                'coordinates': crop_coords,
                                'crop_type': 'overlap',
                                'crop_index': i,
                                'total_crops': len(overlap_crops)
                            })
                    else:
                        # 单次裁剪
                        crop_img, crop_coords = self._single_crop(
                            img, content_bbox, target_size
                        )
                        
                        filename = self._generate_filename(
                            file_info, size_name, 0, 1, crop_coords
                        )
                        
                        output_path = self.output_dir / size_name / filename
                        crop_img.save(output_path, format=self.output_format, quality=self.output_quality)
                        
                        cropped_files.append({
                            'filename': filename,
                            'path': str(output_path),
                            'size_name': size_name,
                            'width': target_width,
                            'height': target_height,
                            'coordinates': crop_coords,
                            'crop_type': 'single',
                            'crop_index': 0,
                            'total_crops': 1
                        })
                
                return [{
                    **file_info,
                    'cropping_status': 'success',
                    'cropped_files': cropped_files
                }]
                
        except Exception as e:
            self.logger.error(f"裁剪单张图片时发生错误: {e}")
            return [{
                **file_info,
                'cropping_status': 'error',
                'error': str(e),
                'cropped_files': []
            }]
    
    async def _crop_single_image(self, file_info: Dict[str, Any]) -> List[Dict[str, Any]]:
        """裁剪单张图片"""
        try:
            with Image.open(file_info['processed_path']) as img:
                content_bbox = file_info['content_bbox'][:4]  # x1, y1, x2, y2
                confidence = file_info.get('confidence', 0.0)
                
                cropped_files = []
                
                # 为每个目标尺寸进行裁剪
                for size_name, target_size in self.target_sizes.items():
                    target_width, target_height = target_size
                    
                    # 检查是否需要重叠裁剪
                    if self._needs_overlap_cropping(content_bbox, target_size):
                        # 重叠裁剪
                        overlap_crops = self._smart_overlap_crop(
                            img, content_bbox, target_size, self.overlap_ratio
                        )
                        
                        for i, (crop_img, crop_coords) in enumerate(overlap_crops):
                            filename = self._generate_filename(
                                file_info, size_name, i, len(overlap_crops), crop_coords
                            )
                            
                            output_path = self.output_dir / size_name / filename
                            crop_img.save(output_path, format=self.output_format, quality=self.output_quality)
                            
                            cropped_files.append({
                                'filename': filename,
                                'path': str(output_path),
                                'size_name': size_name,
                                'width': target_width,
                                'height': target_height,
                                'coordinates': crop_coords,
                                'crop_type': 'overlap',
                                'crop_index': i,
                                'total_crops': len(overlap_crops)
                            })
                    else:
                        # 单次裁剪
                        crop_img, crop_coords = self._single_crop(
                            img, content_bbox, target_size
                        )
                        
                        filename = self._generate_filename(
                            file_info, size_name, 0, 1, crop_coords
                        )
                        
                        output_path = self.output_dir / size_name / filename
                        crop_img.save(output_path, format=self.output_format, quality=self.output_quality)
                        
                        cropped_files.append({
                            'filename': filename,
                            'path': str(output_path),
                            'size_name': size_name,
                            'width': target_width,
                            'height': target_height,
                            'coordinates': crop_coords,
                            'crop_type': 'single',
                            'crop_index': 0,
                            'total_crops': 1
                        })
                
                return [{
                    **file_info,
                    'cropping_status': 'success',
                    'cropped_files': cropped_files
                }]
                
        except Exception as e:
            self.logger.error(f"裁剪单张图片时发生错误: {e}")
            return [{
                **file_info,
                'cropping_status': 'error',
                'error': str(e),
                'cropped_files': []
            }]
    
    def _needs_overlap_cropping(self, content_bbox: List[float], target_size: List[int]) -> bool:
        """判断是否需要重叠裁剪"""
        x1, y1, x2, y2 = content_bbox
        content_width = x2 - x1
        content_height = y2 - y1
        target_width, target_height = target_size
        
        # 如果内容区域任一维度超过目标尺寸，则需要重叠裁剪
        return content_width > target_width or content_height > target_height
    
    def _smart_overlap_crop(self, image: Image.Image, content_bbox: List[float], 
                           target_size: List[int], overlap_ratio: float = 0.25) -> List[Tuple[Image.Image, List[float]]]:
        """
        智能重叠裁剪算法 - 优化版，确保覆盖所有内容
        内存优化：延迟加载和及时释放
        
        Args:
            image: 原始图片
            content_bbox: 内容区域边界框 [x1, y1, x2, y2]
            target_size: 目标尺寸 [width, height]
            overlap_ratio: 重叠比例 (0.2-0.3)
        
        Returns:
            裁剪切片列表 [(裁剪图片, 坐标信息), ...]
        """
        x1, y1, x2, y2 = content_bbox
        content_width = x2 - x1
        content_height = y2 - y1
        target_width, target_height = target_size
        
        self.logger.info(f"开始智能重叠裁剪: 内容区域{content_width}x{content_height}, 目标尺寸{target_width}x{target_height}")
        
        crops = []
        
        # 计算步长（考虑重叠）
        step_x = int(target_width * (1 - overlap_ratio))
        step_y = int(target_height * (1 - overlap_ratio))
        
        # 确保步长至少为1，避免无限循环
        step_x = max(1, step_x)
        step_y = max(1, step_y)
        
        # 计算需要的裁剪数量 - 优化算法确保完整覆盖
        if content_width > target_width:
            num_crops_x = max(1, math.ceil((content_width - target_width) / step_x) + 1)
        else:
            num_crops_x = 1
            
        if content_height > target_height:
            num_crops_y = max(1, math.ceil((content_height - target_height) / step_y) + 1)
        else:
            num_crops_y = 1
        
        self.logger.info(f"计划裁剪网格: {num_crops_x}x{num_crops_y} = {num_crops_x * num_crops_y} 个裁剪")
        
        # 生成裁剪区域 - 优化坐标计算
        for i in range(num_crops_x):
            for j in range(num_crops_y):
                # 计算裁剪区域左上角坐标
                if num_crops_x == 1:
                    # 单个裁剪，居中对齐
                    crop_x = x1 + (content_width - target_width) // 2
                else:
                    # 多个裁剪，按步长均匀分布（保持重叠）
                    crop_x = x1 + i * step_x
                    
                    # 最后一个裁剪特殊处理：确保不超出右边界，但保持重叠
                    if i == num_crops_x - 1 and crop_x + target_width > x2:
                        crop_x = x2 - target_width
                
                if num_crops_y == 1:
                    # 单个裁剪，居中对齐
                    crop_y = y1 + (content_height - target_height) // 2
                else:
                    # 多个裁剪，按步长均匀分布（保持重叠）
                    crop_y = y1 + j * step_y
                    
                    # 最后一个裁剪特殊处理：确保不超出底边界，但保持重叠
                    if j == num_crops_y - 1 and crop_y + target_height > y2:
                        crop_y = y2 - target_height
                
                # 确保不超出内容区域
                crop_x = max(x1, min(crop_x, x2 - target_width))
                crop_y = max(y1, min(crop_y, y2 - target_height))
                
                # 计算裁剪区域右下角坐标
                crop_x2 = crop_x + target_width
                crop_y2 = crop_y + target_height
                
                # 确保不超出图片边界
                crop_x2 = min(crop_x2, image.width)
                crop_y2 = min(crop_y2, image.height)
                
                # 内存优化：直接裁剪，避免不必要的复制
                if (crop_x2 - crop_x) < target_width or (crop_y2 - crop_y) < target_height:
                    crop_img = self._pad_to_target_size(
                        image, (crop_x, crop_y, crop_x2, crop_y2), target_size
                    )
                else:
                    # 使用copy=False避免不必要的内存复制
                    crop_img = image.crop((crop_x, crop_y, crop_x2, crop_y2))
                
                crops.append((crop_img, [crop_x, crop_y, crop_x2, crop_y2]))
                self.logger.debug(f"生成裁剪 {i+1},{j+1}: ({crop_x},{crop_y}) -> ({crop_x2},{crop_y2})")
        
        self.logger.info(f"智能重叠裁剪完成，生成 {len(crops)} 个裁剪")
        return crops
    
    def _single_crop(self, image: Image.Image, content_bbox: List[float], 
                    target_size: List[int]) -> Tuple[Image.Image, List[float]]:
        """单次裁剪"""
        x1, y1, x2, y2 = content_bbox
        target_width, target_height = target_size
        content_width = x2 - x1
        content_height = y2 - y1
        
        # 计算缩放比例
        scale_x = target_width / content_width
        scale_y = target_height / content_height
        scale = min(scale_x, scale_y)  # 使用较小的缩放比例确保内容完整
        
        # 计算缩放后的尺寸
        scaled_width = int(content_width * scale)
        scaled_height = int(content_height * scale)
        
        # 裁剪内容区域
        content_img = image.crop((x1, y1, x2, y2))
        
        # 缩放图片
        scaled_img = content_img.resize((scaled_width, scaled_height), Image.Resampling.LANCZOS)
        
        # 创建目标尺寸的画布（使用白色背景）
        if image.mode == 'RGBA':
            result_img = Image.new('RGB', (target_width, target_height), (255, 255, 255))
        else:
            result_img = Image.new('RGB', (target_width, target_height), (255, 255, 255))
        
        # 计算居中位置
        paste_x = (target_width - scaled_width) // 2
        paste_y = (target_height - scaled_height) // 2
        
        # 粘贴缩放后的图片
        if scaled_img.mode != 'RGB':
            scaled_img = scaled_img.convert('RGB')
        result_img.paste(scaled_img, (paste_x, paste_y))
        
        # 计算在原图中的坐标
        original_coords = [
            x1 + (paste_x / scale),
            y1 + (paste_y / scale),
            x1 + ((paste_x + scaled_width) / scale),
            y1 + ((paste_y + scaled_height) / scale)
        ]
        
        return result_img, original_coords
    
    def _pad_to_target_size(self, image: Image.Image, crop_bbox: List[float], 
                           target_size: List[int]) -> Image.Image:
        """将裁剪区域填充到目标尺寸"""
        x1, y1, x2, y2 = crop_bbox
        target_width, target_height = target_size
        
        # 裁剪区域
        crop_img = image.crop((x1, y1, x2, y2))
        
        # 创建目标尺寸的画布（使用白色背景）
        result_img = Image.new('RGB', (target_width, target_height), (255, 255, 255))
        
        # 计算居中位置
        paste_x = (target_width - crop_img.width) // 2
        paste_y = (target_height - crop_img.height) // 2
        
        # 粘贴裁剪图片
        if crop_img.mode != 'RGB':
            crop_img = crop_img.convert('RGB')
        result_img.paste(crop_img, (paste_x, paste_y))
        
        return result_img
    
    def _generate_filename(self, file_info: Dict[str, Any], size_name: str, 
                          crop_index: int, total_crops: int, coords: List[float]) -> str:
        """生成文件名"""
        original_name = Path(file_info['original_name']).stem
        
        if self.naming_convention == 'intelligent':
            if total_crops == 1:
                # 单次裁剪
                coords_str = f"{int(coords[0])}_{int(coords[1])}_{int(coords[2])}_{int(coords[3])}"
                filename = f"{original_name}_{size_name}_{coords_str}.png"
            else:
                # 重叠裁剪
                filename = f"{original_name}_{size_name}_{crop_index}_{total_crops}.png"
        else:
            # 简单命名
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            if total_crops == 1:
                filename = f"{original_name}_{size_name}_{timestamp}.png"
            else:
                filename = f"{original_name}_{size_name}_{crop_index}_{total_crops}_{timestamp}.png"
        
        return filename
    
    def get_crop_statistics(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """获取裁剪统计信息"""
        total_files = len(results)
        successful_files = sum(1 for r in results if r.get('cropping_status') == 'success')
        total_crops = sum(len(r.get('cropped_files', [])) for r in results)
        
        size_stats = {}
        for result in results:
            for cropped_file in result.get('cropped_files', []):
                size_name = cropped_file['size_name']
                size_stats[size_name] = size_stats.get(size_name, 0) + 1
        
        return {
            'total_files': total_files,
            'successful_files': successful_files,
            'failed_files': total_files - successful_files,
            'total_crops': total_crops,
            'size_distribution': size_stats,
            'average_crops_per_file': total_crops / total_files if total_files > 0 else 0
        }